package reverse_proxy

import (
	"context"
	"errors"
	"fmt"
	"net"
	"runtime"
	"sync/atomic"
	"time"

	"gin-gateway/core/log"
)

var (
	ErrServerClosed     = errors.New("tcp: Server closed")
	ErrAbortHandler     = errors.New("tcp: abort TCPHandler")
	LocalAddrContextKey = &contextKey{"local-addr"}
)

type contextKey struct {
	name string
}

func (ck *contextKey) String() string {
	return "tcp_proxy context value " + ck.name
}

type TCPHandler interface {
	ServeTCP(ctx context.Context, conn net.Conn)
}

type Server struct {
	Addr    string
	Handler TCPHandler
	BaseCtx context.Context

	WriteTimeout     time.Duration
	ReadTimeout      time.Duration
	KeepAliveTimeout time.Duration

	inShutdown  int32
	doneChan    chan struct{}
	TCPListener *net.TCPListener
}

func (s *Server) shuttingDown() bool {
	return atomic.LoadInt32(&s.inShutdown) != 0
}

func (s *Server) ListenAndServe() error {
	if s.shuttingDown() {
		return ErrServerClosed
	}
	if s.doneChan == nil {
		s.doneChan = make(chan struct{})
	}

	addr, err := net.ResolveTCPAddr("tcp", s.Addr)
	if err != nil {
		return errors.New("need addr")
	}
	ln, err := net.ListenTCP("tcp", addr)
	if err != nil {
		return err
	}
	return s.Serve(ln)
}

func (s *Server) Serve(l *net.TCPListener) error {
	s.TCPListener = l
	defer s.TCPListener.Close()

	if s.BaseCtx == nil {
		s.BaseCtx = context.Background()
	}

	for {
		rw, e := l.AcceptTCP()
		if e != nil {
			select {
			case <-s.doneChan:
				return ErrServerClosed
			default:
			}
			fmt.Printf("accept fail, err: %v\n", e)
			continue
		}
		go s.serve(s.BaseCtx, rw)
	}
}

func (s *Server) serve(ctx context.Context, Conn net.Conn) {
	// 设置参数
	if d := s.ReadTimeout; d != 0 {
		err := Conn.SetReadDeadline(time.Now().Add(d))
		if err != nil {
			fmt.Println("serve set read deadline error", err)
			return
		}
	}
	if d := s.WriteTimeout; d != 0 {
		err := Conn.SetWriteDeadline(time.Now().Add(d))
		if err != nil {
			fmt.Println("serve set read deadline error", err)
			return
		}
	}
	if d := s.KeepAliveTimeout; d != 0 {
		// 操作系统保活消息的发送设置
		if tcpConn, ok := Conn.(*net.TCPConn); ok {
			err := tcpConn.SetKeepAlive(true)
			if err != nil {
				fmt.Println("serve set keep alive error", err)
				return
			}
			err = tcpConn.SetKeepAlivePeriod(d)
			if err != nil {
				fmt.Println("serve set keep alive period error", err)
				return
			}
		}
	}

	defer func() {
		if err := recover(); err != nil && err != ErrAbortHandler {
			const size = 64 << 10
			buf := make([]byte, size)
			buf = buf[:runtime.Stack(buf, false)]
			log.Error("tcp: panic serving %v: %v\n%s", Conn.RemoteAddr().String(), err, buf)
		}
		// 回到这里就关闭连接
		Conn.Close()
	}()

	ctx = context.WithValue(ctx, LocalAddrContextKey, Conn.LocalAddr())
	if s.Handler == nil {
		log.Fatal("handler empty")
	}
	s.Handler.ServeTCP(ctx, Conn)
}

func (s *Server) Close() error {
	if atomic.CompareAndSwapInt32(&s.inShutdown, 0, 1) {
		close(s.doneChan)     // 关闭channel
		s.TCPListener.Close() // 执行listener关闭
	}
	return nil
}

func ListenAndServe(addr string, handler TCPHandler) error {
	server := &Server{
		Addr:     addr,
		Handler:  handler,
		doneChan: make(chan struct{}),
	}
	return server.ListenAndServe()
}
